import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import yaml
from models import NightPedestrianPredictor
from data.ncpd_dataset import NCPDDataset
from training.losses import MultiTaskLoss
from training.optimizers import create_optimizer
import logging

def train_model():
    # 加载配置
    with open('configs/model_config.yaml', 'r') as f:
        config = yaml.safe_load(f)
    
    # 设置设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 创建模型
    model = NightPedestrianPredictor(config).to(device)
    
    # 数据集
    train_dataset = NCPDDataset('path/to/train/data', config)
    val_dataset = NCPDDataset('path/to/val/data', config)
    
    train_loader = DataLoader(train_dataset, 
                             batch_size=config['training']['batch_size'],
                             shuffle=True,
                             num_workers=4)
    val_loader = DataLoader(val_dataset,
                           batch_size=config['training']['batch_size'],
                           shuffle=False,
                           num_workers=4)
    
    # 损失函数和优化器
    criterion = MultiTaskLoss()
    optimizer = create_optimizer(model, config)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, 
        T_max=config['training']['epochs']
    )
    
    # 训练循环
    best_val_loss = float('inf')
    
    for epoch in range(config['training']['epochs']):
        model.train()
        train_loss = 0.0
        
        for batch_idx, (infrared, radar, context, labels) in enumerate(train_loader):
            infrared = infrared.to(device)
            radar = radar.to(device)
            context = {k: v.to(device) for k, v in context.items()}
            labels = {k: v.to(device) for k, v in labels.items()}
            
            optimizer.zero_grad()
            
            # 前向传播
            intent_probs, trajectories, _ = model(infrared, radar, context)
            
            # 计算损失
            loss = criterion(intent_probs, trajectories, labels)
            
            # 反向传播
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
            
            if batch_idx % 100 == 0:
                logging.info(f'Epoch: {epoch} | Batch: {batch_idx} | Loss: {loss.item():.4f}')
        
        # 验证
        val_loss = validate_model(model, val_loader, criterion, device)
        
        # 学习率调度
        scheduler.step()
        
        # 保存最佳模型
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), 'best_model.pth')
        
        logging.info(f'Epoch: {epoch} | Train Loss: {train_loss/len(train_loader):.4f} | Val Loss: {val_loss:.4f}')

def validate_model(model, val_loader, criterion, device):
    model.eval()
    total_loss = 0.0
    
    with torch.no_grad():
        for infrared, radar, context, labels in val_loader:
            infrared = infrared.to(device)
            radar = radar.to(device)
            context = {k: v.to(device) for k, v in context.items()}
            labels = {k: v.to(device) for k, v in labels.items()}
            
            intent_probs, trajectories, _ = model(infrared, radar, context)
            loss = criterion(intent_probs, trajectories, labels)
            total_loss += loss.item()
    
    return total_loss / len(val_loader)

if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    train_model()